#include "framework.h"
#include "PluginConfigApi.h"

#include "readerwriterqueue.h"

#include <evr.h>

#include <d3d9.h>
#include <dxva2api.h>

#include <mfapi.h>
#include <mferror.h>
#include <mftransform.h>

#include <future>

class DecodedSample
{
public:
	MFT_OUTPUT_DATA_BUFFER dataBuffer;
	DWORD dwStatus;
	HRESULT result;
};

bool forceSoftwareDecoding;
bool useAsyncDecoding;
int numAllocatedSamples;

IDirect3DDeviceManager9* deviceManager;

IMFTransform* mfTransform;
IMFVideoSampleAllocator* allocator;

moodycamel::BlockingReaderWriterQueue<DecodedSample> decodedSampleQ(30);
std::future<void> futureDecode;

void decodeAllAvailable(DWORD streamID);


HRESULT initialize(IMFTransform* transform)
{
	HRESULT result = S_OK;

	IMFMediaType* type = NULL;

	// Release what we already had beforehand
	SAFE_RELEASE(mfTransform);
	if (allocator)
		ASSERT(allocator->UninitializeSampleAllocator());
	SAFE_RELEASE(allocator);
	while (true)
	{
		DecodedSample frame;
		if (!decodedSampleQ.try_dequeue(frame))
			break;
		SAFE_RELEASE(frame.dataBuffer.pSample);
		SAFE_RELEASE(frame.dataBuffer.pEvents);
	}

	ASSERT(transform->GetOutputCurrentType(0, &type));

	ASSERT(MFCreateVideoSampleAllocator(IID_IMFVideoSampleAllocator, (void**)&allocator));
	ASSERT(allocator->SetDirectXManager(deviceManager));
	ASSERT(allocator->InitializeSampleAllocator(numAllocatedSamples, type));

	mfTransform = transform;
	mfTransform->AddRef();

end:
	SAFE_RELEASE(type);

	return result;
}

VTABLE_HOOK(HRESULT, IMFTransform, ProcessInput, DWORD dwInputStreamID,
	IMFSample* pSample, DWORD dwFlags)
{
	// This method is only hooked if using async decoding
	HRESULT result = S_OK;

	if (mfTransform != This)
		ASSERT(initialize(This));

	// Always wait for any pending ProcessOutput to finish
	if (futureDecode.valid())
		futureDecode.wait();

	ASSERT(originalProcessInput(This, dwInputStreamID, pSample, dwFlags));

	PRINT("[DivaMovie] Start async decoding\n");
	futureDecode = std::async(decodeAllAvailable, dwInputStreamID);
end:
	return result;
}

VTABLE_HOOK(HRESULT, IMFTransform, ProcessOutput, DWORD dwFlags, DWORD cOutputBufferCount,
	MFT_OUTPUT_DATA_BUFFER* pOutputSamples, DWORD* pdwStatus)
{
	HRESULT result = S_OK;

	if (mfTransform != This)
		ASSERT(initialize(This));

	if (useAsyncDecoding)
	{
		PRINT("[DivaMovie] async ProcessOutput\n");
		DecodedSample frame;
		decodedSampleQ.wait_dequeue(frame);
		PRINT("[DivaMovie] pulled decode result from queue\n");
		pOutputSamples[0] = frame.dataBuffer;
		*pdwStatus = frame.dwStatus;
		result = frame.result;
	}
	else
	{
		if (pOutputSamples->pSample != NULL || deviceManager == NULL || cOutputBufferCount != 1)
			return originalProcessOutput(This, dwFlags, cOutputBufferCount, pOutputSamples, pdwStatus);

		ASSERT(allocator->AllocateSample(&(pOutputSamples->pSample)));
		ASSERT(originalProcessOutput(This, dwFlags, cOutputBufferCount, pOutputSamples, pdwStatus));
	}
end:
	return result;
}

void hookDecodeMethods(IMFTransform* mft)
{
	if (useAsyncDecoding)
	{
		INSTALL_VTABLE_HOOK(mft, ProcessInput, 24);
	}
	INSTALL_VTABLE_HOOK(mft, ProcessOutput, 25);
}

void decodeAllAvailable(DWORD streamID)
{
	HRESULT result;

	while (true)
	{
		DecodedSample sampleToQ;

		sampleToQ.dwStatus = 0;
		sampleToQ.dataBuffer.dwStreamID = streamID;
		ASSERT(allocator->AllocateSample(&(sampleToQ.dataBuffer.pSample)));
		sampleToQ.dataBuffer.dwStatus = 0;
		sampleToQ.dataBuffer.pEvents = NULL;
		sampleToQ.result = originalProcessOutput(mfTransform, 0, 1, &(sampleToQ.dataBuffer), &(sampleToQ.dwStatus));
		if (SUCCEEDED(sampleToQ.result))
		{
			PRINT("[DivaMovie] async originalProcessOutput succeeded\n");
		}
		else
		{
			PRINT("[DivaMovie] async originalProcessOutput failed (error code: 0x%08x)\n", (unsigned int)sampleToQ.result);
		}
		decodedSampleQ.enqueue(sampleToQ);
		if (FAILED(sampleToQ.result))
			break;
	}
end:
	return;
}

VTABLE_HOOK(HRESULT, IMFTransform, ProcessMessage, MFT_MESSAGE_TYPE eMessage, ULONG_PTR ulParam)
{
	if (forceSoftwareDecoding)
	{
		if (eMessage == MFT_MESSAGE_SET_D3D_MANAGER) {
			PRINT("[DivaMovie] Force Software Decoding enabled\n");
			hookDecodeMethods(This);
			return S_OK;
		}
		return originalProcessMessage(This, eMessage, ulParam);
	}
	else
	{
		HRESULT result = originalProcessMessage(This, eMessage, ulParam);
		if (eMessage == MFT_MESSAGE_SET_D3D_MANAGER && result == MF_E_UNSUPPORTED_D3D_TYPE)
		{
			PRINT("[DivaMovie] This system does not support DXVA hardware decoding\n")
			result = originalProcessMessage(This, eMessage, NULL);
			if (SUCCEEDED(result))
				hookDecodeMethods(This);
		}
		return result;
	}
	
}

HOOK(void*, IMFTransformInitializer, 0x140420B90, void* a1, void* a2, IMFTransform** transform)
{
	void* result = originalIMFTransformInitializer(a1, a2, transform);
	INSTALL_VTABLE_HOOK(*transform, ProcessMessage, 23);
	return result;
}

HOOK(HRESULT, DXVA2CreateDirect3DDeviceManager, PROC_ADDRESS("dxva2.dll", "DXVA2CreateDirect3DDeviceManager9"),
	UINT* pResetToken, IDirect3DDeviceManager9** ppDeviceManager)
{
	HRESULT result = originalDXVA2CreateDirect3DDeviceManager(pResetToken, ppDeviceManager);

	if (FAILED(result))
		return result;

	if (deviceManager)
		deviceManager->Release();

	deviceManager = *ppDeviceManager;
	deviceManager->AddRef();

	return result;
}

void loadConfig() {
	forceSoftwareDecoding = GetPrivateProfileIntW(L"general", L"force_software_decoding", 0, CONFIG_FILE) > 0 ? true : false;
	useAsyncDecoding = GetPrivateProfileIntW(L"general", L"async_decoding", 0, CONFIG_FILE) > 0 ? true : false;
	numAllocatedSamples = GetPrivateProfileIntW(L"general", L"num_allocated_samples", 3, CONFIG_FILE);
}

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
	if (ul_reason_for_call == DLL_PROCESS_ATTACH)
	{
		loadConfig();
		INSTALL_HOOK(DXVA2CreateDirect3DDeviceManager);
		INSTALL_HOOK(IMFTransformInitializer);
	}

	return TRUE;
}

PluginConfig::PluginConfigOption config[] = {
	{ PluginConfig::CONFIG_BOOLEAN, new PluginConfig::PluginConfigBooleanData{
		L"force_software_decoding", L"general", CONFIG_FILE, L"Force Software Decoding",
		L"Use software decoding even on systems that support DXVA hardware decoding.", false } },
	{ PluginConfig::CONFIG_BOOLEAN, new PluginConfig::PluginConfigBooleanData{
		L"async_decoding", L"general", CONFIG_FILE, L"Use Async Decoding",
		L"If software decoding is enabled, decode in a separate thread.", false } },
	{ PluginConfig::CONFIG_NUMERIC, new PluginConfig::PluginConfigNumericData{
		L"num_allocated_samples", L"general", CONFIG_FILE, L"Allocate Samples:",
		L"Sets the number of video samples the Allocator shoud allocate.\nWhen Use Async Decoding is disabled, setting this to 3 suffices;\nWhen Use Async Decoding is enabled, set it to a larger number such as 100.",
		3, 1, 600 } },
};

extern "C" __declspec(dllexport) LPCWSTR GetPluginName(void)
{
	return L"DivaMovie";
}

extern "C" __declspec(dllexport) LPCWSTR GetPluginDescription(void)
{
	return L"DivaMovie Plugin by Skyth and twistero\n\nDivaMovie enables movies on systems that does not support DXVA hardware decoding.";
}

extern "C" __declspec(dllexport) PluginConfig::PluginConfigArray GetPluginOptions(void)
{
	return PluginConfig::PluginConfigArray{ _countof(config), config };
}